# OCR相关路由
import base64
import json
from typing import List, Optional, Literal
from fastapi import APIRouter, UploadFile, File, HTTPException, Body, Form
from fastapi.responses import StreamingResponse
from starlette.responses import JSONResponse, HTMLResponse

from app.config.settings import settings
from app.utils.logger import get_logger
from app.core.ocr_processor import OCRProcessor

logger = get_logger(__name__)
router = APIRouter(tags=["OCR处理"])
# router = APIRouter(prefix="/ocr", tags=["OCR处理"])
ocr_processor = OCRProcessor()


# 从文件读取描述内容
def load_description_from_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        return "描述文件未找到"


# 动态描述装饰器
def dynamic_description(file_path):
    def decorator(func):
        # 为函数添加属性，用于标识描述文件路径
        func._description_file = file_path
        return func

    return decorator


# 文本识别API - 文件上传版本
@router.post("/ocr/text/file",
             summary="文本OCR识别（文件上传）",
             description="通过上传图像文件进行文本识别")
async def text_ocr_file(
        images: List[UploadFile] = File(
            ...,
            description="要识别的图像文件列表"
        ),
        language: Literal["ch", "en", "fr", "german", "korean", "japan"] = Body(
            "ch",
            description="识别语言，支持: ch(简体中文)、en(英文)、fr(法文)、german(德文)、korean(韩文)、japan(日文)",
        ),
        detail: bool = Body(
            False,
            description="是否返回详细信息，包括边界框和置信度"
        )
):
    """文本OCR接口 - 文件上传版本"""
    logger.info(f"收到文本OCR文件上传请求, language: {language}, detail: {detail}")
    lang = settings.get('ocr.lang')

    # 检查语言参数 - 增加对 lang 为 None 的检查
    if lang is None or language not in lang:
        logger.error(f"无效的语言参数: {language} 或未配置支持的语言列表")
        raise HTTPException(status_code=400, detail="无效的语言参数或服务未正确配置")

    try:
        # 处理上传的文件
        image_contents = []

        # 处理上传的文件
        if images:
            for image in images:
                # 检查是否是有效的上传文件（跳过空文件）
                if hasattr(image, 'filename') and image.filename:
                    contents = await image.read()
                    image_contents.append(contents)

        if not image_contents:
            raise HTTPException(status_code=400, detail="未提供任何图像数据")

        # 异步处理OCR
        results = []
        for content in image_contents:
            result = await ocr_processor.process_text_recognition(content, language, detail)
            results.append(result)

        # logger.info("文本OCR文件上传处理成功")
        # 如果只有一张图片，保持原有返回格式
        if len(results) == 1:
            return {"msg": "ok", "data": results[0]}
        else:
            return {"msg": "ok", "data": results}
    except Exception as e:
        logger.error(f"文本OCR文件上传处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 文本识别API - Base64版本
@router.post("/ocr/text/base64",
             summary="文本OCR识别（Base64）")
@dynamic_description("descriptions/text_ocr_base64.md")
async def text_ocr_base64(
        image_base64_strings: List[str] = Form(
            ...,
            description="base64编码的图像字符串（可以提交多个）"
        ),
        language: Literal["ch", "en", "fr", "german", "korean", "japan"] = Body(
            "ch",
            description="识别语言，支持: ch(简体中文)、en(英文)、fr(法文)、german(德文)、korean(韩文)、japan(日文)",
        ),
        detail: bool = Form(
            False,
            description="是否返回详细信息，包括边界框和置信度"
        )
):
    """文本OCR接口 - Base64版本"""
    logger.info(f"收到文本OCR Base64请求, language: {language}, detail: {detail}")
    lang = settings.get('ocr.lang')

    # 检查语言参数 - 增加对 lang 为 None 的检查
    if lang is None or language not in lang:
        logger.error(f"无效的语言参数: {language} 或未配置支持的语言列表")
        raise HTTPException(status_code=400, detail="无效的语言参数或服务未正确配置")

    try:
        # 处理base64字符串
        image_contents = []

        for base64_string in image_base64_strings:
            try:
                # 移除可能的数据URL前缀
                if base64_string.startswith('data:image'):
                    base64_string = base64_string.split(',', 1)[1]
                image_data = base64.b64decode(base64_string)
                image_contents.append(image_data)
            except Exception as e:
                logger.error(f"Base64解码失败: {str(e)}")
                raise HTTPException(status_code=400, detail=f"Base64解码失败: {str(e)}")

        if not image_contents:
            raise HTTPException(status_code=400, detail="未提供任何图像数据")

        # 异步处理OCR
        results = []
        for content in image_contents:
            result = await ocr_processor.process_text_recognition(content, language, detail)
            results.append(result)

        logger.info("文本OCR Base64处理成功")
        # 如果只有一张图片，保持原有返回格式
        if len(results) == 1:
            return JSONResponse(content={"msg": "ok", "data": results[0]})
        else:
            return JSONResponse(content={"msg": "ok", "data": results})
    except Exception as e:
        logger.error(f"文本OCR Base64处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 版面分析API
@router.post("/ocr/layout", summary="版面分析")
@dynamic_description("descriptions/layout_analysis.md")
async def layout_analysis(
        image: UploadFile = File(...),
        detect_tables: bool = True,
        detect_formulas: bool = True
):
    """版面分析接口"""
    logger.info(f"收到版面分析请求, detect_tables: {detect_tables}, detect_formulas: {detect_formulas}")

    try:
        contents = await image.read()
        # 异步处理版面分析
        result = await ocr_processor.process_layout_analysis(contents, detect_tables, detect_formulas)
        logger.info("版面分析处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"版面分析处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 表格结构识别API
@router.post("/ocr/table", summary="表格结构识别")
@dynamic_description("descriptions/table_structure_recognition.md")
async def table_recognition(
        image: UploadFile = File(...),
        structure_only: bool = False
):
    """表格结构识别接口"""
    logger.info(f"收到表格结构识别请求, structure_only: {structure_only}")

    try:
        contents = await image.read()
        # 异步处理表格结构识别
        result = await ocr_processor.process_table_structure_recognition(contents, structure_only)
        logger.info("表格结构识别处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"表格结构识别处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 公式识别API
@router.post("/ocr/formula", summary="公式识别")
@dynamic_description("descriptions/formula_recognition.md")
async def formula_recognition(
        image: UploadFile = File(...),
        format_type: str = "latex"
):
    """公式识别接口"""
    logger.info(f"收到公式识别请求, format_type: {format_type}")

    try:
        contents = await image.read()
        # 异步处理公式识别
        result = await ocr_processor.process_formula_recognition(contents, format_type)
        logger.info("公式识别处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"公式识别处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 文档方向分类API
@router.post("/ocr/doc_orientation", summary="文档方向分类")
@dynamic_description("descriptions/doc_orientation.md")
async def doc_orientation(image: UploadFile = File(...)):
    """文档方向分类接口"""
    logger.info("收到文档方向分类请求")

    try:
        contents = await image.read()
        # 异步处理文档方向分类
        result = await ocr_processor.process_doc_orientation(contents)
        logger.info("文档方向分类处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"文档方向分类处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 印章文本检测API
@router.post("/ocr/seal_text", summary="印章文本检测")
@dynamic_description("descriptions/seal_text_detection.md")
async def seal_text_detection(image: UploadFile = File(...)):
    """印章文本检测接口"""
    logger.info("收到印章文本检测请求")

    try:
        contents = await image.read()
        # 异步处理印章文本检测
        result = await ocr_processor.process_seal_text_detection(contents)
        logger.info("印章文本检测处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"印章文本检测处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 图表解析API
@router.post("/ocr/chart", summary="图表解析")
@dynamic_description("descriptions/chart_parsing.md")
async def chart_parsing(
        image: UploadFile = File(...),
        prompt: str = ""
):
    """图表解析接口"""
    logger.info(f"收到图表解析请求, prompt: {prompt}")

    try:
        contents = await image.read()
        # 异步处理图表解析
        result = await ocr_processor.process_chart_parsing(contents, prompt)
        logger.info("图表解析处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"图表解析处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 视觉语言模型API
@router.post("/ocr/vlm", summary="视觉语言模型分析")
@dynamic_description("descriptions/vlm.md")
async def vision_language_analysis(
        image: UploadFile = File(...),
        prompt: str = "Describe this image"
):
    """视觉语言模型分析接口"""
    logger.info(f"收到视觉语言模型分析请求, prompt: {prompt}")

    try:
        contents = await image.read()
        # 异步处理视觉语言模型分析
        result = await ocr_processor.process_vision_language_analysis(contents, prompt)
        logger.info("视觉语言模型分析处理成功")
        return {"msg": "ok", "data": result}
    except Exception as e:
        logger.error(f"视觉语言模型分析处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
